package edu.fsu.cs.buzbee.steggo;

import java.awt.Point;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;

import javax.imageio.ImageIO;

/**
 * Matrix Encoder hides data by using a Matrix Coding 
 * @author Ben
 *
 */
public class MatrixEncoder {
	long offset = 0;
	BufferedImage image;
	public MatrixEncoder(BufferedImage base)
	{
		
		// Argument sanity check before accepting construction
		if (base == null)
			throw new NullPointerException("Base image cannot be null");
		if (base.getHeight() == 0 || base.getWidth() == 0)
			throw new RuntimeException("Base image has a dimension value 0");
		
		image = base;

			
	}
	
	Point linearToPoint(long offset)
	{
		long y = offset / image.getWidth();
		long x = offset - (image.getWidth()*y);
		
		return new Point((int)x, (int)y);
	}
	
	public void encode(File in, File out) throws IOException
	{
		if (in == null || out == null || !in.exists())
			throw new RuntimeException("Invalid input or output file");
		if (!in.canRead())
			throw new IOException("No access to read input file");
		if (!out.exists() && !out.createNewFile())
			throw new IOException("Cannot create output file");
		if (!out.canWrite())
			throw new IOException("No access to write output file");
	
		long filesize = in.length() * 8;
		
		long k = chooseK(filesize);
		if (k == 0)
		{
			throw new IOException("Message will not fit in file.");
		}
		long n = (1L << k) - 1L;
		
		long componentBytes = (image.getWidth() * image.getHeight()*3);
		long numN = (int)Math.floor(componentBytes / n);
		
		// Sanity check - will the message fit
		
		if ((n * numN) > componentBytes)
			{
				throw new IOException("The input file will not fit in the "+componentBytes+" component bytes with k="+k+", n = "+n+", numN = "+numN);
			}
		
		System.out.println("MatrixEncoder: Encoding "+filesize +" bits in blocks of k="+k);
		System.out.println("\t\tinto "+componentBytes + " component byte LSBs in "+numN+" blocks of n="+n);
	
		
		FileInputStream fin = new FileInputStream(in);
		long size = 0;
		for (int i = 0; i < numN; ++i)
		{
			// For each N sized cover block (bits n*i to (n*i+n - 1))
			// We encode K bits of the message block (from i*k to i*k+k-1)
			// By changing a single bit 

			long MB = 0;
			// Get k bits, designate them MB
			for (long k_this = k; k_this > 0; k_this -= 8)
			{
				int nextByte = fin.read();
				if (nextByte == -1)
					nextByte = 0;
				else
					++size;
				
				MB |= (nextByte << (k_this - 8));
			
			}
			
			long alpha = MB ^ b(i, n);
			
			if (alpha != 0)
			{
				// Flip bit i*n + alpha -1
				setLSB(i * n + alpha -1, 1 - getLSB(i * n + alpha - 1));
			}
			
		}
		fin.close();
		
		
		ImageIO.write(image, "png", out);
		
		System.out.println("Done. REMEMBER K="+k+" AND SIZE="+size+", this is necessary for decoding.");
		
	
		
		
	}
	
	/**
	 * Calculates the b(cb) function for the cover block of size N starting at bit N*cb_offset
	 * First cb, cb_offset = 0
	 * Second cb, cb_offset = 1
	 * ...
	 * @param cb_offset
	 * @return
	 */
	public long b(int cb_offset, long n)
	{
		long result = 0;
		for (long position = cb_offset*n, j = 1; j <= n; j++, position++)
		{
			result ^= getLSB(position) * j;
		}
		return result;
	}
	
	/**
	 * Selects the size of message blocks K to be the maximum that will allow this message to fit into N = 2^k -1 bit cover blocks
	 * Returns 0 on error
	 * @param message_size
	 * @return
	 */
	public int chooseK(long message_size)
	{

		
		// Choose largest K such that message can be divided into K size blocks
		// and message can be split into N=2^k-1 size blocks
		long c =  image.getWidth() * image.getHeight() * 3; // Number of cover blocks
		for (int k = 56; k > 0; k-=8)
		{
			// 2^k -1
			long n = (1L << k) - 1L;
			
			
			// If the cover bits cannot be split into even one message of size n, this k is unsuitable
			if ((c / n) < 1)
				continue;
			
			// How many groups of size K do we need to encode the entire message?
			long groupsOfK = (long)Math.ceil((double)message_size / k);
			
			// If having this many groups of K would require more groups of size N than we we have cover bits, this k wont work
			if (n*groupsOfK > c)
				continue;
			
			return k;
			
		
		}
		return 0;
	}
	
	/**
	 * Gets least significant bit at location
	 * Example: Get the LSB of the 2nd (green) component of the first pixel with location = 1
	 *          Get the LSB of the 3rd component of the 5th pixel with location = 17
	 * @param location
	 * @return
	 */
	public int getLSB(long location)
	{
		if (location < 0 || location > (image.getWidth() * image.getHeight() * 3))
				throw new IndexOutOfBoundsException("Location out of bounds: "+location);
		
		long pixel = location / 3;
		int mask;
		if (location % 3 == 0)
			mask = 0x00010000;
		else if (location % 3 == 1)
			mask = 0x00000100;
		else
			mask = 0x00000001;
		
		Point p = linearToPoint(pixel);
		int rgb = image.getRGB(p.x, p.y);
		if ((rgb & mask) != 0) return 1;
		else return 0;
	}
	/**
	 * Sets least significant bit at location to given 0 <= value <= 1
	 * Example: Set the LSB of the 2nd (Green) component of the first pixel with location = 1
	 *          Set the LSB of the 3rd component of the 5th pixel with location = 17
	 * @param location
	 * @param value
	 */
	public void setLSB(long location, int value)
	{
		if (value < 0 || value > 1)
			throw new IndexOutOfBoundsException("Value must be 0 or 1");
		if (location < 0 || location > (image.getWidth() * image.getHeight() * 3))
			throw new IndexOutOfBoundsException("Location out of bounds: "+location);
		
		long pixel = location / 3;
		
		
		Point p = linearToPoint(pixel);
		int rgb = image.getRGB(p.x, p.y);
		
		int mask = 0;
		if (value == 0)
		{
			if (location % 3 == 0)
				mask = 0xFFFEFFFF;
			else if (location % 3 == 1)
				mask = 0xFFFFFEFF;
			else if (location % 3 == 2)
				mask = 0xFFFFFFFE;
			
			image.setRGB(p.x, p.y, rgb & mask);
		}
		else if (value == 1)
		{
			if (location % 3 == 0)
				mask = 0x00010000;
			else if (location % 3 == 1)
				mask = 0x00000100;
			else if (location % 3 == 2)
				mask = 0x00000001;
			image.setRGB(p.x, p.y, rgb | mask);
		}
		
	}
}
